import torch
from math import pi, sqrt

"""
Saleh-Valenzuela 模型
    1. 用于模拟无线信道的多径传播特性。
    2. 信道模型包括随机路径增益、路径时延和方向角信息。
    3. 适用于 毫米波通信 和 大规模 MIMO 系统。
该信道矩阵可用于仿真无线通信中的波束成形、链路级性能分析等场景。
"""

def Hermitian(matrix):
    """计算复数矩阵的共轭转置。"""
    return matrix.transpose(-1, -2).conj()

def kron(a, b):
    """
    计算两个矩阵的克罗内克积 (Kronecker Product)。
    输入:
        a: 维度为 [batch_size, N]
        b: 维度为 [batch_size, N]
    输出:
        c: 克罗内克积的结果，维度为 [batch_size, N * N]
    """
    batch_size = a.shape[0]  # 获取 batch_size 大小
    a = a.reshape(batch_size, -1, 1)  # 将 a 扩展为 [batch_size, N, 1]
    b = b.reshape(batch_size, 1, -1)  # 将 b 扩展为 [batch_size, 1, N]
    c = a @ b  # 计算矩阵乘法，对应 Kronecker 积
    return c.reshape(batch_size, -1)  # 拉平为 [batch_size, N * N]

def kron_add(a,b):
    #a与b维度为[batch_size,N],输出应为[batch_size,N*N]
    batch_size = a.shape[0]
    a = a.reshape(batch_size,-1,1)
    b = b.reshape(batch_size,1,-1)
    c = a + b
    return c.reshape(batch_size,-1) #输出的维度a在前，b在后

def DFT_matrix(N):
    """
    生成 N 阶离散傅里叶变换 (DFT) 矩阵。
    输入:
        N: 矩阵的阶数
    输出:
        W: DFT 矩阵，维度为 [N, N]
    """
    n = torch.arange(N).reshape(N, 1).cuda()  # 生成行向量 n
    k = torch.arange(N).reshape(1, N).cuda()  # 生成列向量 k
    W = torch.exp(-1j * 2 * pi * n * k / N) / sqrt(N)  # 根据公式生成 DFT 矩阵
    return W

class SVChannelModel:
    def __init__(self, batch_size, Nc, paths, Nt, sigma_2_alpha):
        self.batch_size = batch_size
        self.Nc = Nc
        self.paths = paths
        self.Nt = Nt
        self.sigma_2_alpha = sigma_2_alpha
    
    def build(self):
        return SV_Channel(self.batch_size, self.Nc, self.paths, self.Nt, self.sigma_2_alpha)
        
def SV_Channel(batch_size, Nc, paths, Nt, sigma_2_alpha):
    """
    生成 Saleh-Valenzuela (SV) 信道模型。
    输入:
        batch_size: 批大小
        Nc: 载波数
        N: 多径数量
        Nt: 发射端天线数，格式为 [Nx, Ny]
        sigma_2_alpha: 每条路径的能量
    输出:
        H_f: 频域信道矩阵，维度为 [batch_size, Nc, N_t]
    """
    gpu = torch.cuda.current_device()
    d = 0.5  # 天线间距，单位波长

    # 生成路径时延，范围 [0, Tau)，Tau = Nc 表示最大路径时延
    Tau = Nc
    tau = torch.rand(batch_size, 1, paths).cuda(gpu) * Tau

    # 定义发射端和接收端天线数量
    N_r = 1  # 接收端天线数
    N_t = Nt[0] * Nt[1]  # 发射端天线总数

    # 随机生成路径角度 (发射/接收端)
    fait = torch.rand(batch_size, 1, paths).cuda(gpu) * 2 * pi  # 发射端方位角
    fair = torch.rand(batch_size, 1, paths).cuda(gpu) * 2 * pi  # 接收端方位角
    theatt = torch.rand(batch_size, 1, paths).cuda(gpu) * 2 * pi  # 发射端仰角
    theatr = torch.rand(batch_size, 1, paths).cuda(gpu) * 2 * pi  # 接收端仰角

    # 初始化发射/接收端阵列响应和路径增益
    A_t = torch.zeros(batch_size, N_t, paths).cuda(gpu) + 0j  # 发射端阵列响应
    alpha = torch.zeros(batch_size, 1, paths).cuda(gpu) + 0j  # 多径路径增益
    n_t1 = torch.arange(Nt[0]).reshape(1, Nt[0], 1).cuda(gpu)  # 发射端天线第一个维度索引
    n_t2 = torch.arange(Nt[1]).reshape(1, Nt[1], 1).cuda(gpu)  # 发射端天线第二个维度索引
    A_r = torch.ones(batch_size, 1, paths).cuda(gpu) + 0j  # 接收端阵列响应 (单天线)
    H_f = torch.zeros(batch_size, Nc, N_t).cuda(gpu) + 0j  # 初始化频域信道矩阵

    # 遍历多径路径，计算信道参数
    for i in range(paths):
        # 发射端阵列响应计算
        at1 = torch.exp(-2j * pi * d * n_t1 * torch.cos(fait[:, :, i:(i + 1)]) * torch.sin(theatt[:, :, i:(i + 1)]))
        at2 = torch.exp(-2j * pi * d * n_t2 * torch.sin(fait[:, :, i:(i + 1)]))
        A_t[:, :, i] = kron(at1, at2)  # 克罗内克积得到总的发射端阵列响应

        # 路径增益 (复数高斯分布，方差为 sigma_2_alpha / 2)
        alpha[:, :, i] = (torch.randn(batch_size, 1).cuda(gpu) + 1j * torch.randn(batch_size, 1).cuda(gpu)) * sqrt(sigma_2_alpha / 2)

    # 遍历载波，计算频域信道
    for k in range(Nc):
        # 每个路径的贡献，考虑路径增益和时延
        P = alpha * torch.exp(-1j * 2 * pi * tau * k / Nc)
        P = torch.diag_embed(P).reshape(batch_size, paths, paths)  # 将路径贡献转为对角矩阵

        # 计算频域信道矩阵
        H_f[:, k:(k + 1), :] = torch.matmul(torch.matmul(A_r, P), Hermitian(A_t))

    return H_f


def LAMP_GMMV(Nc, Nt, SNR_dB, K, L, pho, paths, H):
    """
    基于 LAMP 的广义多用户多径信道恢复。
    输入:
        Nc: 载波数
        Nt: 发射天线数
        snr: 信噪比
        K: 用户数
        L: 导频时隙数
        pho: 采样密度，决定感知矩阵过采样程度
        H: 原始信道矩阵，形状为 [batch_size, Nc, Nt]
    输出:
        H_hat: 恢复的信道矩阵，形状为 [batch_size, Nc, Nt]
    """
    gpu = torch.cuda.current_device()
    # 初始化常量
    factor = 100 * L / 32  # 归一化因子，用于信号幅度调整
    # B = 4096               # 数据块大小（未用到）
    # batch_size = 64             # 每次处理的批次大小
    # Step = 3000            # 迭代步数（未用到）
    T = 10                 # LAMP 网络的迭代次数（未定义，需替换为具体值）

    # 从参数列表提取系统参数
    snr = snr = 10**(SNR_dB / 10)
    sigma = 1 / snr         # 噪声功率
    
    # 构造离散傅里叶变换矩阵 A_t，用于信道估计
    N_2 = pho * 8
    A = torch.zeros(Nt, N_2 * N_2).cuda(gpu) + 0j
    for i in range(N_2):
        for j in range(N_2):
            # 构建 Kronecker 相位数组
            a_h = torch.range(0, 8 - 1).reshape(1, 8).cuda(gpu) / 8 / pho * i * 2 * pi
            a_v = torch.range(0, 8 - 1).reshape(1, 8).cuda(gpu) / 8 / pho * j * 2 * pi
            a = kron_add(a_h, a_v).reshape(Nt)  # 使用自定义函数 kron_add 生成阵列响应
            A[:, i * N_2 + j] = torch.exp(1j * a)  # 构造复数矩阵
    A_t = A  # DFT 矩阵

    with torch.no_grad():  # 禁用梯度计算（推理阶段）

        # 随机生成导频信号
        S_digital = torch.randn(L, Nc, K, 1).cuda(gpu) + 1j * torch.randn(L, Nc, K, 1).cuda(gpu)
        S_analog = torch.rand(L, 1, Nt, 4).cuda(gpu) * 2 * pi

        # 加载预训练 LAMP 模型和导频矩阵
        net = torch.load('./pretrain/channel/GMMV_LAMP' + str(T) + '_' + str(L) + 'pilots_' +
                         str(paths) + 'paths_' + str(SNR_dB) + 'SNR' + '.pth')
        net.eval()  # 设置网络为推理模式
        S_digital = torch.load('./pretrain/channel/S_digital' + str(L) + 'pilots_' +
                               str(paths) + 'paths_' + str(SNR_dB) + 'SNR' + '.pth')
        S_analog = torch.load('./pretrain/channel/S_analog' + str(L) + 'pilots_' +
                              str(paths) + 'paths_' + str(SNR_dB) + 'SNR' + '.pth')

        # 初始化观测信号和感知矩阵
        Y = torch.zeros(H.shape[0], Nc, L).cuda(gpu) + 1j  # [batch_size, Nc, L]
        Phi = torch.zeros(Nc, L, Nt).cuda(gpu) + 0j       # [Nc, L, Nt]

        # 遍历时隙，生成观测信号 Y 和感知矩阵 Phi
        for l in range(L):
            # 构建导频矩阵 s
            s = (torch.exp(1j * S_analog[l, :, :, :]) @ S_digital[l, :, :, :]).reshape(Nc, Nt, 1)

            # 对导频归一化
            s_sigma = torch.sqrt(torch.sum(torch.abs(s * s)))
            s = s / s_sigma.reshape(1, 1, 1) * torch.min(s_sigma, torch.tensor(sqrt(Nc))).reshape(1, 1, 1)

            # 加入噪声
            n = (torch.randn(H.shape[0], Nc) + 1j * torch.randn(H.shape[0], Nc)).cuda(gpu) * sqrt(sigma / 2)

            # 构建观测信号
            a = s.permute(0, 2, 1)  # 调整维度以匹配
            y = (a @ H.reshape(-1, Nc, Nt, 1)).reshape(-1, Nc) + n  # [batch_size, Nc]
            Y[:, :, l] = y
            Phi[:, l, :] = a[:, 0, :]

        # 计算感知矩阵 Theta
        Theta = Phi @ A_t
        Theta = Theta / sqrt(factor)
        Y = Y / sqrt(factor)

        # 使用 LAMP 网络恢复信道
        H_t_hat, V = net(Y, Nc, Nt * pho * pho, L, Theta)  # 恢复的信道 [batch_size, 过采样, Nc]
        H_t_hat = H_t_hat.permute(0, 2, 1).reshape(Y.shape[0], Nc, Nt * pho * pho, 1)
        H_hat = A_t @ H_t_hat  # 投影到实际天线
        H_hat = H_hat.reshape(-1, Nc, Nt)  # 恢复信道矩阵 [batch_size, Nc, Nt]

    return H_hat